import torch
from sklearn.model_selection import KFold

import os
import json
import os.path

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.patches import Polygon

from Pretrain.models.ssl_head import SSLHead

from glob import glob

plt.rcParams['font.family'] = 'Times New Roman'
plt.rcParams['font.weight'] = 'bold'
plt.rcParams['mathtext.fontset'] = 'stix'

def draw_plot(data):
    cls = data.keys()

    fig, ax = plt.subplots()

    for i, label in enumerate(cls):
        ax.plot(data[label], label=label)

    ax.legend()

    plt.ylabel('Loss', rotation=90, fontdict={'size': 16})
    plt.xlabel('Epoch')
    plt.show()


def SplitDataToJson(datalist, target='./'):
    """

    :param datalist: 传入的 json 文件地址
    :param target: 生成的 json 文件的目标地址
    :param fold: K 折交叉验证
    :param key: 文件名称
    :return:
    """
    with open(datalist) as f:
        json_data = json.load(f)

    print(json_data.keys())
    new_json = {'modality': json_data['modality'],
                'labels': json_data['labels'],
                'training': [],
                'test': []}

    training_list = []
    test_list = []

    KF = KFold(5, shuffle=True, random_state=28)

    for i, (train_index, val_index) in enumerate(KF.split(json_data['training'])):
        for index in val_index:
            data_split = {'image': None, 'label': None, 'fold': i}
            data = json_data['training'][index]
            data_split['image'] = data['image'][2:]
            data_split['label'] = data['label'][2:]
            training_list.append(data_split)

    for path in json_data['test']:
        path = path[2:]
        test_list.append(path)

    new_json['training'] = training_list
    new_json['test'] = test_list

    with open(os.path.join(target, 'Amos22/jsons/amos_t1.json'), 'w') as f:
        json.dump(new_json, f, indent=4)


def visual_data():
    import SimpleITK as sitk
    import matplotlib.pyplot as plt

    luna = '/Users/qlc/Desktop/Dataset/Luna16/1.3.6.1.4.1.14519.5.2.1.6279.6001.109002525524522225658609808059.mhd/1.3.6.1.4.1.14519.5.2.1.6279.6001.109002525524522225658609808059.mhd'
    brats = '/Users/qlc/Desktop/Dataset/brats2021/TrainingData/BraTS2021_01647/BraTS2021_01647_flair.nii.gz'
    lung = '/Users/qlc/Desktop/Dataset/MSD/Task06_Lung/imagesTr/lung_001.nii.gz'
    amos = '/Users/qlc/Desktop/Dataset/AMOS22/imagesTr/amos_0001.nii.gz'
    brats_label = '/Users/qlc/Desktop/Dataset/Brats2021/TrainingData/BraTS2021_00000/BraTS2021_00000_seg.nii.gz'

    img = sitk.ReadImage(brats)
    img = sitk.GetArrayFromImage(img)

    img = img[90, 45:205, 45:205]
    fig, ax = plt.subplots()
    ax.imshow(img, cmap='gray')
    ax.set_xticks([])
    ax.set_yticks([])
    plt.show()
    plt.close(fig)

    cnt = 0
    for i in range(8):
        for j in range(8):
            patch = img[20 * i: 20 * (i + 1), 20 * j:20 * (j + 1)]
    
            fig, ax = plt.subplots()
            ax.imshow(patch, cmap='gray')
            ax.set_xticks([])
            ax.set_yticks([])
            # plt.show()
            plt.savefig(f'/Users/qlc/Desktop/src/{cnt}.png')
            plt.close(fig)
            # break
            cnt+= 1


if __name__ == '__main__':
    import SimpleITK as sitk
    
    # path = '/Users/qlc/Desktop/Dataset/Brats2021/TrainingData/BraTS2021_00005/BraTS2021_00005_seg.nii.gz'
    # path = '/Users/qlc/Desktop/Dataset/AMOS22/labelsTr/amos_0005.nii.gz'
    # luna = '/Users/qlc/Desktop/Dataset/Luna16/1.3.6.1.4.1.14519.5.2.1.6279.6001.109002525524522225658609808059.mhd/1.3.6.1.4.1.14519.5.2.1.6279.6001.109002525524522225658609808059.mhd'
    # brats = '/Users/qlc/Desktop/Dataset/brats2021/TrainingData/BraTS2021_01647/BraTS2021_01647_flair.nii.gz'
    # lung = '/Users/qlc/Desktop/Dataset/MSD/Task06_Lung/imagesTr/lung_001.nii.gz'
    # amos = '/Users/qlc/Desktop/Dataset/AMOS22/imagesTr/amos_0001.nii.gz'
    brats_label = '/Users/qlc/Desktop/Dataset/Brats2023/Adult_Glioma/ASNR-MICCAI-BraTS2023-GLI-Challenge-TrainingData/BraTS-GLI-00000-000/BraTS-GLI-00000-000-seg.nii.gz'
    
    img = sitk.ReadImage(brats_label)
    seg = sitk.GetArrayFromImage(img)
    print(seg.shape, np.unique(seg))
    cmap = plt.cm.Paired
    # cmap = plt.cm.Pastel1

    d, _, _ = seg.shape
    # 将类别对应到指定的颜色，如果需要修改透明背景时，将每一个颜色的最后一维设置为 0
    seg_ = np.zeros((seg.shape[0], 
                     seg.shape[1],
                     seg.shape[2],
                     4)) 
    # print(seg)
    # for i, j in enumerate([0, 2, 1, 4]):
    #     seg_[seg==j] = cmap(i)
    # print(seg_.shape)
    # print(np.unique(seg_[1]))

    # seg_[seg==0] = cmap(0)
    # for i in range(12):
    #     seg_[seg==i] = cmap(i)
        
    # cmap = plt.cm.Pastel1

    # seg_[seg==12] = cmap(2)

    # seg_[seg==13] = cmap(7)
    
    # seg_[seg==14] = cmap(5)
    
    # seg_[seg==15] = cmap(6)

    # seg = seg_
    
    # plt.figure()
    # plt.imshow(seg_[60, :, :])
    # plt.show()

    # for i in range(0, d, 2):
    #     seg_name = os.path.join('/Users/qlc/Desktop/test/b5', f'seg_{i}.png')
    #     plt.imsave(seg_name, seg[i, :, :], dpi=150)
    
    # datalist = '/Users/qlc/Desktop/a.json'
    
    # with open(datalist) as f:
    #     json_data = json.load(f)
    
    # print(json_data['training'][0])
    
    # data_root = '/Users/qlc/Desktop/Dataset/Brats2020/MICCAI_BraTS2020_TrainingData'
    
    # path_list = glob(data_root + '/*')
    # # print(path_list)
    
    # targe_json = {'training': []}
    
    # for i in path_list:
    #     if 'csv' not in i:
    #         case_info = {'image': [],
    #                     'label': None}
    #         for j in sorted(glob(i + '/*.nii.gz')):
    #             if 'seg' in j:
    #                 case_info['label'] = j
    #             else:
    #                 case_info['image'].append(j)

    #         targe_json['training'].append(case_info)

    # with open('/Users/qlc/Desktop/a.json', 'w') as f:
    #     json.dump(targe_json, f)


    # training_list = []
    # test_list = []

    # KF = KFold(10, shuffle=True, random_state=28)

    # for i, (train_index, val_index) in enumerate(KF.split(json_data['training'])):
    #     for index in val_index:
    #         data_split = {'image': None, 'label': None, 'fold': i}
    #         data = json_data['training'][index]
    #         data_split['image'] = data['image'][2:]
    #         data_split['label'] = data['label'][2:]
    #         training_list.append(data_split)
    # target_json = {'training': []}
    # target_json['training'] = training_list


    # with open('/Users/qlc/Desktop/b.json', 'w') as f:
    #     json.dump(target_json, f)



    # path = '/Users/qlc/Desktop/gmim/amos/finetune/gmim_5_5_3_5_1/Amos2022Dataset_0/Amos2022Dataset_0_ckpt_best.pt'
    
    # model = torch.load(path, map_location='cpu')
    # print(model.keys())
    # print(model['epoch'], model['best_acc'])




